"""
Author:VictorChen
Date:02/10/2022
"""
import cv2
import numpy as np
import math


class Seam:
    def __init__(self, img):
        self.h, self.w = 0, 0
        self.energy_mat = 0
        self.img = 0
        self.renew(img)

    def renew(self, img):
        self.img = img
        h, w = img.shape
        img1 = img.astype(np.float64)
        energy_mat = np.zeros((h, w))

        """
        Using Sobel operator to get the energy matrix
        """

        for row in range(h):
            for col in range(w):
                if (row == 0) | (row == h - 1) | (col == 0) | (col == w - 1):
                    energy_mat[row, col] = 255
                else:
                    dx = sum(img1[row - 1:row + 2, col - 1] - img1[row - 1:row + 2, col + 1]) + img1[row, col - 1] - \
                         img1[
                             row, col + 1]
                    dy = sum(img1[row - 1, col - 1:col + 2] - img1[row + 1, col - 1:col + 2]) + img1[row - 1, col] - \
                         img1[
                             row + 1, col]
                    energy_mat[row, col] = math.sqrt(dx * dx + dy * dy)
        energy_mat[0] = energy_mat[1]
        energy_mat[h - 1] = energy_mat[h - 2]

        energy_mat[:, 0] = energy_mat[:, 1]
        energy_mat[:, w - 1] = energy_mat[:, w - 2]

        self.h, self.w = h, w
        self.energy_mat = energy_mat.astype(np.uint8)

    def width_find_seam(self):
        path = np.zeros((self.h, self.w)).astype(np.int32)
        path_long = np.zeros((self.h, self.w))
        path_long[0] = self.energy_mat[0]

        """
        Using DP algorithm to find the shortest path from bottom to up
        """
        for i in range(1, self.h):
            v = np.append(2256, path_long[i - 1], 2256)
            for j in range(self.w):
                idx = np.argmin(v[j:j + 3])

                path[i, j] = idx - 1 + j
                # if(path[i,j]<0):
                #     print(v[j:j+3])

                # Accumulate the length
                # print(int(path[i, j]))
                path_long[i, j] = path_long[i - 1, int(path[i, j])] + self.energy_mat[i, j]

        destination = np.argmin(path_long[self.h - 1])
        seamlist = [destination]
        for i in range(1, self.h):
            destination = path[self.h - i, destination]
            seamlist.append(destination)
        seamlist.reverse()
        return seamlist

    def cut_width(self):
        lst = self.width_find_seam()
        new_img = np.zeros((self.h, self.w - 1))
        for i in range(self.h):
            for j in range(lst[i]):
                new_img[i, j] = self.img[i, j]
            for j in range(lst[i] + 1, self.w):
                new_img[i, j - 1] = self.img[i, j]
        self.renew(new_img)
        return new_img

    def cut_height(self):
        self.renew(self.img.T)
        new_img= self.cut_width().T
        self.renew(new_img)
        return new_img

    def add_width(self):
        lst = self.width_find_seam()
        new_img = np.zeros((self.h, self.w + 1))
        for i in range(self.h):
            for j in range(lst[i]+1):
                new_img[i, j] = self.img[i, j]
            j = lst[i]+1
            new_img[i , j] = self.img[i , j]*0.5+self.img[i,j+1]*0.5
            for j in range(lst[i] + 1, self.w):
                new_img[i, j + 1] = self.img[i, j]
        self.renew(new_img)
        return new_img

    def show(self):
        lst = self.width_find_seam()
        new_img = np.zeros((self.h, self.w + 1))
        for i in range(self.h):
            for j in range(lst[i] + 1):
                new_img[i, j] = self.img[i, j]
            j = lst[i] + 1
            new_img[i, j] = 0
            for j in range(lst[i] + 1, self.w):
                new_img[i, j + 1] = self.img[i, j]
        self.renew(new_img)
        return new_img
    
    def add_height(self):
        self.renew(self.img.T)
        new_img= self.add_width().T
        self.renew(new_img)
        return new_img
        

if __name__ == "__main__":
    image = cv2.imread("demo/grey.jpg")
    # print(image)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    cv2.imwrite("demo/grey.jpg",image)

    """
    test cut_height
    """
    image1 = Seam(image)
    for i in range(50):
        image1.cut_height()
    cv2.imwrite("demo/cut_height.jpg", image1.img)
    
    """
    test cut_width
    """
    image2 = Seam(image)
    for i in range(50):
        image2.cut_width()
    cv2.imwrite("demo/cut_width.jpg", image2.img)
    
    """
    test add_height
    """
    image2 = Seam(image)
    for i in range(10):
        image2.add_height()
    cv2.imwrite("demo/add_height.jpg", image2.img)
    
    """
    test add_width
    """
    image2 = Seam(image)
    for i in range(50):
        image2.add_width()
    cv2.imwrite("demo/add_width.jpg", image2.img)

    """
    track the paths to show the effect
    """

    image2 = Seam(image)
    for i in range(50):
        image2.show()
    cv2.imwrite("demo/show.jpg", image2.img)

